"""
任务规划器

这个模块实现了 Agent 的任务规划能力，展示了 AI Agent 如何：
1. 分解复杂任务为简单步骤
2. 制定执行计划和时间表
3. 管理任务依赖关系
4. 跟踪执行进度
5. 动态调整计划

任务规划是处理复杂问题的关键能力。
"""

import asyncio
import json
from typing import List, Dict, Any, Optional, Set
from dataclasses import dataclass, field
from enum import Enum
from datetime import datetime, timedelta


class TaskStatus(Enum):
    """任务状态枚举"""
    PENDING = "pending"         # 待执行
    IN_PROGRESS = "in_progress" # 执行中
    COMPLETED = "completed"     # 已完成
    FAILED = "failed"          # 失败
    CANCELLED = "cancelled"    # 已取消
    BLOCKED = "blocked"        # 被阻塞


class TaskPriority(Enum):
    """任务优先级枚举"""
    LOW = 1
    MEDIUM = 2
    HIGH = 3
    URGENT = 4


@dataclass
class Task:
    """任务类"""
    id: str                                    # 任务ID
    title: str                                 # 任务标题
    description: str                           # 任务描述
    status: TaskStatus = TaskStatus.PENDING    # 任务状态
    priority: TaskPriority = TaskPriority.MEDIUM  # 优先级
    
    # 时间相关
    created_at: datetime = field(default_factory=datetime.now)
    estimated_duration: Optional[timedelta] = None
    actual_start_time: Optional[datetime] = None
    actual_end_time: Optional[datetime] = None
    deadline: Optional[datetime] = None
    
    # 依赖关系
    dependencies: Set[str] = field(default_factory=set)  # 依赖的任务ID
    dependents: Set[str] = field(default_factory=set)    # 依赖此任务的任务ID
    
    # 执行相关
    required_tools: List[str] = field(default_factory=list)  # 需要的工具
    resources: Dict[str, Any] = field(default_factory=dict)  # 需要的资源
    constraints: List[str] = field(default_factory=list)     # 约束条件
    
    # 结果和反馈
    result: Optional[Any] = None               # 执行结果
    error_message: Optional[str] = None        # 错误信息
    progress: float = 0.0                      # 进度 (0-1)
    
    # 元数据
    metadata: Dict[str, Any] = field(default_factory=dict)
    tags: List[str] = field(default_factory=list)


@dataclass
class ExecutionPlan:
    """执行计划"""
    plan_id: str
    tasks: List[Task]
    total_estimated_time: timedelta
    critical_path: List[str]  # 关键路径上的任务ID
    created_at: datetime = field(default_factory=datetime.now)


class TaskPlanner:
    """
    任务规划器类
    
    负责将复杂的用户请求分解为可执行的任务序列，
    并管理任务的执行过程。这展示了 AI Agent 如何
    处理复杂的多步骤问题。
    """
    
    def __init__(self):
        self.tasks: Dict[str, Task] = {}
        self.plans: Dict[str, ExecutionPlan] = {}
        self.current_plan_id: Optional[str] = None
        
    async def create_plan(self, user_request: str, context: Dict[str, Any] = None) -> ExecutionPlan:
        """
        创建执行计划
        
        这个方法展示了完整的任务规划过程：
        1. 分析用户请求
        2. 分解为子任务
        3. 确定任务依赖
        4. 估算时间和资源
        5. 生成执行计划
        """
        
        # 1. 分析请求复杂度
        complexity_analysis = self._analyze_complexity(user_request)
        
        # 2. 分解任务
        tasks = await self._decompose_request(user_request, complexity_analysis, context or {})
        
        # 3. 建立依赖关系
        self._establish_dependencies(tasks)
        
        # 4. 优化任务顺序
        optimized_tasks = self._optimize_task_order(tasks)
        
        # 5. 估算总时间
        total_time = self._estimate_total_time(optimized_tasks)
        
        # 6. 计算关键路径
        critical_path = self._calculate_critical_path(optimized_tasks)
        
        # 7. 创建执行计划
        plan_id = f"plan_{datetime.now().strftime('%Y%m%d_%H%M%S')}"
        plan = ExecutionPlan(
            plan_id=plan_id,
            tasks=optimized_tasks,
            total_estimated_time=total_time,
            critical_path=critical_path
        )
        
        # 8. 保存计划和任务
        self.plans[plan_id] = plan
        for task in optimized_tasks:
            self.tasks[task.id] = task
        
        self.current_plan_id = plan_id
        
        return plan
    
    def _analyze_complexity(self, request: str) -> Dict[str, Any]:
        """分析请求复杂度"""
        request_lower = request.lower()
        
        # 复杂度指标
        complexity_indicators = {
            "multiple_actions": len([w for w in ["和", "然后", "接着", "以及", "同时"] if w in request]),
            "conditional_logic": len([w for w in ["如果", "否则", "当", "while", "if"] if w in request_lower]),
            "data_processing": len([w for w in ["分析", "处理", "计算", "统计", "排序"] if w in request]),
            "external_resources": len([w for w in ["搜索", "下载", "获取", "查询"] if w in request]),
            "creative_tasks": len([w for w in ["设计", "创建", "生成", "写", "画"] if w in request]),
        }
        
        total_complexity = sum(complexity_indicators.values())
        
        return {
            "total_score": total_complexity,
            "level": "simple" if total_complexity <= 2 else "medium" if total_complexity <= 5 else "complex",
            "indicators": complexity_indicators,
            "estimated_tasks": max(1, total_complexity // 2 + 1)
        }
    
    async def _decompose_request(self, request: str, complexity: Dict, context: Dict) -> List[Task]:
        """分解用户请求为具体任务"""
        tasks = []
        request_lower = request.lower()
        
        # 基础任务ID生成器
        task_counter = 1
        
        def create_task(title: str, description: str, tools: List[str] = None, 
                       priority: TaskPriority = TaskPriority.MEDIUM,
                       estimated_minutes: int = 5) -> Task:
            nonlocal task_counter
            task_id = f"task_{datetime.now().strftime('%Y%m%d')}_{task_counter:03d}"
            task_counter += 1
            
            return Task(
                id=task_id,
                title=title,
                description=description,
                priority=priority,
                required_tools=tools or [],
                estimated_duration=timedelta(minutes=estimated_minutes)
            )
        
        # 1. 预处理任务（如果需要）
        if any(keyword in request_lower for keyword in ["搜索", "查找", "获取信息"]):
            tasks.append(create_task(
                "信息收集",
                "搜索和收集相关信息",
                ["search"],
                TaskPriority.HIGH,
                10
            ))
        
        # 2. 分析常见任务模式
        if "计算" in request or any(op in request for op in ["+", "-", "*", "/"]):
            tasks.append(create_task(
                "数学计算",
                "执行数学运算",
                ["calculator"],
                TaskPriority.MEDIUM,
                3
            ))
        
        if any(keyword in request_lower for keyword in ["文件", "保存", "写入", "创建"]):
            tasks.append(create_task(
                "文件操作",
                "执行文件相关操作",
                ["filesystem"],
                TaskPriority.MEDIUM,
                8
            ))
        
        if any(keyword in request_lower for keyword in ["分析", "解释", "说明"]):
            tasks.append(create_task(
                "内容分析",
                "分析和解释内容",
                [],
                TaskPriority.HIGH,
                15
            ))
        
        if any(keyword in request_lower for keyword in ["生成", "创建", "写", "设计"]):
            tasks.append(create_task(
                "内容生成",
                "生成或创建新内容",
                [],
                TaskPriority.MEDIUM,
                20
            ))
        
        # 3. 复杂请求的额外分解
        if complexity["level"] == "complex":
            # 添加规划任务
            tasks.insert(0, create_task(
                "任务规划",
                "制定详细的执行策略",
                [],
                TaskPriority.URGENT,
                5
            ))
            
            # 添加验证任务
            tasks.append(create_task(
                "结果验证",
                "验证执行结果的正确性",
                [],
                TaskPriority.HIGH,
                8
            ))
        
        # 4. 如果没有识别出具体任务，创建通用任务
        if not tasks:
            tasks.append(create_task(
                "处理用户请求",
                f"处理请求: {request}",
                [],
                TaskPriority.MEDIUM,
                10
            ))
        
        # 5. 总是添加总结任务
        if len(tasks) > 1:
            tasks.append(create_task(
                "结果总结",
                "整理和总结执行结果",
                [],
                TaskPriority.LOW,
                5
            ))
        
        return tasks
    
    def _establish_dependencies(self, tasks: List[Task]):
        """建立任务依赖关系"""
        if len(tasks) <= 1:
            return
        
        # 简单的线性依赖：除第一个任务外，每个任务都依赖前一个任务
        for i in range(1, len(tasks)):
            current_task = tasks[i]
            previous_task = tasks[i-1]
            
            current_task.dependencies.add(previous_task.id)
            previous_task.dependents.add(current_task.id)
        
        # 特殊依赖规则
        task_titles = {task.title: task for task in tasks}
        
        # 验证任务依赖所有主要任务
        if "结果验证" in task_titles:
            validation_task = task_titles["结果验证"]
            for task in tasks:
                if task.title not in ["结果验证", "结果总结", "任务规划"]:
                    validation_task.dependencies.add(task.id)
                    task.dependents.add(validation_task.id)
        
        # 总结任务依赖所有其他任务
        if "结果总结" in task_titles:
            summary_task = task_titles["结果总结"]
            for task in tasks:
                if task.title != "结果总结":
                    summary_task.dependencies.add(task.id)
                    task.dependents.add(summary_task.id)
    
    def _optimize_task_order(self, tasks: List[Task]) -> List[Task]:
        """优化任务执行顺序"""
        # 简单的拓扑排序
        task_map = {task.id: task for task in tasks}
        result = []
        visited = set()
        temp_visited = set()
        
        def visit(task_id: str):
            if task_id in temp_visited:
                # 检测到循环依赖
                return
            if task_id in visited:
                return
            
            temp_visited.add(task_id)
            task = task_map[task_id]
            
            # 先访问所有依赖
            for dep_id in task.dependencies:
                if dep_id in task_map:
                    visit(dep_id)
            
            temp_visited.remove(task_id)
            visited.add(task_id)
            result.append(task)
        
        # 访问所有任务
        for task in tasks:
            if task.id not in visited:
                visit(task.id)
        
        return result
    
    def _estimate_total_time(self, tasks: List[Task]) -> timedelta:
        """估算总执行时间"""
        # 考虑并行执行的可能性
        total_time = timedelta()
        
        # 简化：假设任务串行执行
        for task in tasks:
            if task.estimated_duration:
                total_time += task.estimated_duration
        
        # 添加一些缓冲时间
        buffer_time = timedelta(minutes=len(tasks) * 2)
        return total_time + buffer_time
    
    def _calculate_critical_path(self, tasks: List[Task]) -> List[str]:
        """计算关键路径"""
        # 简化实现：找到从开始到结束的最长路径
        if not tasks:
            return []
        
        # 找到没有依赖的起始任务
        start_tasks = [task for task in tasks if not task.dependencies]
        if not start_tasks:
            return [tasks[0].id]  # 如果有循环依赖，返回第一个任务
        
        # 找到没有被依赖的结束任务
        end_tasks = [task for task in tasks if not task.dependents]
        if not end_tasks:
            return [tasks[-1].id]  # 如果有循环依赖，返回最后一个任务
        
        # 简化：返回从第一个开始任务到第一个结束任务的路径
        critical_path = []
        current_task = start_tasks[0]
        visited = set()
        
        while current_task and current_task.id not in visited:
            critical_path.append(current_task.id)
            visited.add(current_task.id)
            
            # 找到下一个任务（优先级最高的依赖者）
            next_task = None
            highest_priority = 0
            
            for dep_id in current_task.dependents:
                dep_task = self.tasks.get(dep_id)
                if dep_task and dep_task.priority.value > highest_priority:
                    highest_priority = dep_task.priority.value
                    next_task = dep_task
            
            current_task = next_task
        
        return critical_path
    
    async def execute_plan(self, plan_id: str = None) -> Dict[str, Any]:
        """
        执行计划
        
        这个方法展示了任务执行的完整流程：
        1. 检查执行条件
        2. 按顺序执行任务
        3. 处理依赖关系
        4. 错误处理和恢复
        5. 进度跟踪
        """
        
        if plan_id is None:
            plan_id = self.current_plan_id
        
        if plan_id not in self.plans:
            return {"error": "计划不存在"}
        
        plan = self.plans[plan_id]
        execution_log = []
        failed_tasks = []
        
        # 按依赖顺序执行任务
        for task in plan.tasks:
            try:
                # 检查依赖是否已完成
                if not self._check_dependencies_completed(task):
                    task.status = TaskStatus.BLOCKED
                    execution_log.append(f"任务 {task.title} 被阻塞：依赖未完成")
                    continue
                
                # 开始执行任务
                task.status = TaskStatus.IN_PROGRESS
                task.actual_start_time = datetime.now()
                
                execution_log.append(f"开始执行任务: {task.title}")
                
                # 模拟任务执行（实际实现中会调用相应的工具）
                result = await self._execute_task(task)
                
                if result.get("success", True):
                    task.status = TaskStatus.COMPLETED
                    task.result = result.get("result")
                    task.progress = 1.0
                    execution_log.append(f"任务完成: {task.title}")
                else:
                    task.status = TaskStatus.FAILED
                    task.error_message = result.get("error", "未知错误")
                    failed_tasks.append(task.id)
                    execution_log.append(f"任务失败: {task.title} - {task.error_message}")
                
                task.actual_end_time = datetime.now()
                
            except Exception as e:
                task.status = TaskStatus.FAILED
                task.error_message = str(e)
                task.actual_end_time = datetime.now()
                failed_tasks.append(task.id)
                execution_log.append(f"任务异常: {task.title} - {str(e)}")
        
        # 生成执行报告
        completed_tasks = [t for t in plan.tasks if t.status == TaskStatus.COMPLETED]
        total_time = sum([
            (t.actual_end_time - t.actual_start_time).total_seconds() 
            for t in completed_tasks 
            if t.actual_start_time and t.actual_end_time
        ], 0)
        
        return {
            "plan_id": plan_id,
            "total_tasks": len(plan.tasks),
            "completed_tasks": len(completed_tasks),
            "failed_tasks": len(failed_tasks),
            "execution_time_seconds": total_time,
            "success_rate": len(completed_tasks) / len(plan.tasks) if plan.tasks else 0,
            "execution_log": execution_log,
            "failed_task_ids": failed_tasks
        }
    
    def _check_dependencies_completed(self, task: Task) -> bool:
        """检查任务依赖是否已完成"""
        for dep_id in task.dependencies:
            dep_task = self.tasks.get(dep_id)
            if not dep_task or dep_task.status != TaskStatus.COMPLETED:
                return False
        return True
    
    async def _execute_task(self, task: Task) -> Dict[str, Any]:
        """执行单个任务（模拟）"""
        # 模拟任务执行时间
        if task.estimated_duration:
            # 实际项目中这里会调用相应的工具和执行逻辑
            await asyncio.sleep(min(1, task.estimated_duration.total_seconds() / 60))  # 缩短模拟时间
        
        # 模拟不同类型任务的执行结果
        if "搜索" in task.title or "信息收集" in task.title:
            return {"success": True, "result": "已收集相关信息"}
        elif "计算" in task.title:
            return {"success": True, "result": "计算完成"}
        elif "文件" in task.title:
            return {"success": True, "result": "文件操作完成"}
        elif "分析" in task.title:
            return {"success": True, "result": "分析报告已生成"}
        else:
            return {"success": True, "result": f"任务 {task.title} 执行完成"}
    
    def get_plan_status(self, plan_id: str = None) -> Dict[str, Any]:
        """获取计划执行状态"""
        if plan_id is None:
            plan_id = self.current_plan_id
        
        if plan_id not in self.plans:
            return {"error": "计划不存在"}
        
        plan = self.plans[plan_id]
        
        status_counts = {}
        for status in TaskStatus:
            status_counts[status.value] = sum(1 for task in plan.tasks if task.status == status)
        
        total_progress = sum(task.progress for task in plan.tasks) / len(plan.tasks) if plan.tasks else 0
        
        return {
            "plan_id": plan_id,
            "total_tasks": len(plan.tasks),
            "status_breakdown": status_counts,
            "overall_progress": total_progress,
            "critical_path": plan.critical_path,
            "estimated_total_time": str(plan.total_estimated_time),
            "created_at": plan.created_at.isoformat()
        }
    
    def get_task_details(self, task_id: str) -> Optional[Dict[str, Any]]:
        """获取任务详细信息"""
        task = self.tasks.get(task_id)
        if not task:
            return None
        
        return {
            "id": task.id,
            "title": task.title,
            "description": task.description,
            "status": task.status.value,
            "priority": task.priority.value,
            "progress": task.progress,
            "created_at": task.created_at.isoformat(),
            "estimated_duration": str(task.estimated_duration) if task.estimated_duration else None,
            "actual_start_time": task.actual_start_time.isoformat() if task.actual_start_time else None,
            "actual_end_time": task.actual_end_time.isoformat() if task.actual_end_time else None,
            "dependencies": list(task.dependencies),
            "dependents": list(task.dependents),
            "required_tools": task.required_tools,
            "result": task.result,
            "error_message": task.error_message
        }